import logging
import time
import uuid

from fastapi import FastAPI, BackgroundTasks, HTTPException

from cache import SessionManager
from models import RetrieveResponse, RetrieveRequest, UpdateRequest, UpdateResponse, BatchUpdateRequest
from services import FaissManager
from services import Neo4jManager

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

neo4j_manager = Neo4jManager()
faiss_manager = FaissManager()
session_manager = SessionManager(neo4j_manager, faiss_manager)

app = FastAPI(title="MemTree Memory System")


@app.post("/v1/memory/sessions/{session_id}/retrieve", response_model=RetrieveResponse)
async def retrieve_memory(session_id: uuid.UUID, request: RetrieveRequest):
    context = await session_manager.retrieve(session_id, request.input, request.k)
    return RetrieveResponse(retrieved_context=context)


@app.post("/v1/memory/sessions/{session_id}/update", response_model=UpdateResponse)
async def update_memory(session_id: uuid.UUID, request: UpdateRequest):  # 移除 background_tasks
    new_node_id = await session_manager.update(
        session_id, request.full_turn_text  # 移除 background_tasks
    )
    return UpdateResponse(new_node_id=new_node_id)


@app.post("/v1/memory/sessions/{session_id}/batch_update")
async def batch_update_memory(session_id: uuid.UUID, request: BatchUpdateRequest):
    start_time = time.time()
    num_texts = len(request.texts)
    logger.info(f"Received batch_update request for session {session_id} with {num_texts} texts.")

    status = await session_manager.batch_update(session_id, request.texts)

    end_time = time.time()
    duration = end_time - start_time
    logger.info(f"Finished batch_update for session {session_id}. Total time: {duration:.2f} seconds.")

    return {"status": "completed", "processed_items": num_texts}


@app.on_event("shutdown")
def shutdown_event():
    neo4j_manager.close()